import logging
from mcp.server.fastmcp import FastMCP
import math
import numpy as np
from scipy.linalg import lu, qr
from typing import List, Tuple
import asyncio

# 日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 创建 MCP 服务实例
mcp = FastMCP("Math Server")

# ========= 基础数学工具 =========
@mcp.tool()
def add(a: float, b: float) -> float:
    return a + b

@mcp.tool()
def subtract(a: float, b: float) -> float:
    return a - b

@mcp.tool()
def multiply(a: float, b: float) -> float:
    return a * b

@mcp.tool()
def divide(a: float, b: float) -> float:
    if b == 0:
        raise ValueError("除数不能为零")
    return a / b

@mcp.tool()
def square_root(a: float) -> float:
    if a < 0:
        raise ValueError("不能计算负数的平方根")
    return math.sqrt(a)

@mcp.tool()
def power(a: float, b: float) -> float:
    return math.pow(a, b)

@mcp.tool()
def absolute(a: float) -> float:
    return abs(a)

# ========= 线性代数工具 =========
@mcp.tool()
def gaussian_elimination(matrix: List[List[float]], vector: List[float]) -> List[float]:
    try:
        A = np.array(matrix, dtype=float)
        b = np.array(vector, dtype=float)
        x = np.linalg.solve(A, b)
        return x.tolist()
    except Exception as e:
        raise ValueError(f"高斯消元法求解失败: {str(e)}")

@mcp.tool()
def lu_decomposition(matrix: List[List[float]]) -> Tuple[List[List[float]], List[List[float]], List[List[float]]]:
    try:
        A = np.array(matrix, dtype=float)
        P, L, U = lu(A)
        return P.tolist(), L.tolist(), U.tolist()
    except Exception as e:
        raise ValueError(f"LU 分解失败: {str(e)}")

@mcp.tool()
def qr_decomposition(matrix: List[List[float]]) -> Tuple[List[List[float]], List[List[float]]]:
    try:
        A = np.array(matrix, dtype=float)
        Q, R = qr(A)
        return Q.tolist(), R.tolist()
    except Exception as e:
        raise ValueError(f"QR 分解失败: {str(e)}")

# ========= 打印已注册工具 =========
def print_registered_tools():
    try:
        tools = asyncio.run(mcp.list_tools())
        logger.info("已注册的工具函数列表:")
        for tool in tools:
            logger.info(f"  -> {tool.name}")
    except Exception as e:
        logger.warning(f"无法列出工具列表: {str(e)}")

# ========= 启动 MCP =========
if __name__ == "__main__":
    logger.info("正在启动 MCP 服务...")
    print_registered_tools()
    try:
        mcp.run(transport="stdio")
        logger.info("MCP 服务已成功启动")
    except Exception as e:
        logger.error(f"MCP 服务启动失败: {str(e)}")
